import torch.nn as nn
from torch.nn.utils import prune
import torch

class Pruner():
    def __init__(self, amount):
        self.amount = amount
        self.pruned = False
    
    def prune_model(self, model):
        self.pruned = True

        for name, module in model.named_modules():
            if isinstance(module, (nn.Conv2d, nn.Linear)):
                prune.l1_unstructured(module, name='weight', amount=self.amount)
                print(f"Pruned {name} with amount {self.amount:.4f}")
    
    def recover(self, model):
        if self.pruned==False:
            return 
        for name, module in model.named_modules():
            if isinstance(module, nn.Conv2d):
                prune.remove(module, name='weight')
            if isinstance(module, nn.Linear):
                prune.remove(module, name='weight')
            self.pruned = False
            
            
class PruningAwareCompressor():
    def __init__(self):
        self.mask_map = {}
        self.update_count_map = {}
        self.stable_map = {}
        
    def compress(self, bucket):
        tensor = bucket.buffer()
        shape = tensor.size()
        tensor = tensor.flatten()
        index = bucket.index()

        if index in self.stable_map and self.stable_map[index]:
            compressed_data = tensor[~self.mask_map[index]]
        else:
            compressed_data = tensor
        return compressed_data, shape
    
    def zero_mask_track(self, bucket, decompressed_tensor):
        index = bucket.index()
        current_mask = (decompressed_tensor == 0).type(torch.bool)
        if index not in self.mask_map or self.mask_map[index].size() != current_mask.size():
            self.mask_map[index] = current_mask
            return 
        previous_mask = self.mask_map[index]
        self.mask_map[index] = previous_mask & current_mask
        
        num_updates = (self.mask_map[index] != previous_mask).sum().item()
        self.stable_track(index=index, num_updates=num_updates)
    
    def stable_track(self, index, num_updates):
        if index not in self.stable_map:
            self.stable_map[index] = False
        if self.stable_map[index]:
            return
        if num_updates == 0:
            if index not in self.update_count_map:
                self.update_count_map[index] = 1
            else:
                self.update_count_map[index] += 1
        else:
            self.update_count_map[index] = 0

        if self.update_count_map[index] >= 10:
            self.stable_map[index] = True
            print(f'index {index} is stable')

    def decompress(self, compressed_data, shape, bucket):
        index = bucket.index()
        
        if compressed_data.size() != shape:
            decompressed_data = torch.zeros(shape, dtype=torch.float32, device=compressed_data.device)
            decompressed_data[~self.mask_map[index]] = compressed_data.type(torch.float32)
            return decompressed_data
        else:
            return compressed_data